import base64
from ctypes import CDLL, c_ubyte, c_char_p
import hashlib
import httplib
import json

# FHMQV key sizes
import os

FHMQV_ECP_PRIVATE_KEY_LENGTH = 64
FHMQV_ECP_PUBLIC_KEY_LENGTH = 129
FHMQV_ECP_SHARED_KEY_LENGTH = 64

IV_SIZE = 16 #128bit IV
ATAG_SIZE = 16 #128bit ATAG
SESSION_ID_SIZE = 8


# server's static public key (this must be known beforehand)
server_static_public_key = "0446b74da517f8b4e2d795beb170298d13ea3b576639efdc0a16f80e771a16a597de7b99764be47d6a5e93825a852501c15ab6d6c979e57f5d2bc7b4a2c9c859eea0cfa6c99effe59cffd542a0070379fa96fda58fb80163e2ddd3157a9b791abd".decode('hex')


# client's own static key-pair (this key-pair is known beforehand)
client_static_private_key = "c4687ad2c37cd34521d8ded1e6e95fcfae234308dda70aa52f9602b663b196fd04b23f58bfc04f40865b50e89aaf82f2".decode('hex')
client_static_public_key = "04e41886cc931a6cfeaabee9a6c322aaa44e40fae249c0b04613985bf18ce092c1056922293493fac594c3eecbbaa64a564b4472b8f09fe6af728338e01c5cbac9f30b4882fee4a7b2e9ae3f4035ef762e766241cc862742e303a07fd265209175".decode('hex')


def hexstr(s):
    return ' '.join(c.encode('hex') for c in s)

def main():
    libgksl = CDLL("libgksl.so")
    conn = httplib.HTTPConnection('localhost', 8181)

    #############################################################################################################
    ### CREATE SESSION

    shared_key = bytearray(FHMQV_ECP_SHARED_KEY_LENGTH)
    client_ephemeral_public_key = bytearray(FHMQV_ECP_PUBLIC_KEY_LENGTH)
    client_ephemeral_private_key = bytearray(FHMQV_ECP_PRIVATE_KEY_LENGTH)

    #generate ephemeral key-pair
    libgksl.fhmqv_ecp_gen_keypair((c_ubyte * FHMQV_ECP_PRIVATE_KEY_LENGTH).from_buffer(client_ephemeral_private_key),
                                  (c_ubyte * FHMQV_ECP_PUBLIC_KEY_LENGTH).from_buffer(client_ephemeral_public_key))

    request = {'id': 0,
               'method': 'create_session',
               'params': {
                   'eph_pub_key': base64.b64encode(client_ephemeral_public_key)
                }}

    conn.request('POST', '/', json.dumps(request))

    resp = conn.getresponse()
    if resp.status != httplib.OK:
        print('server responded with http error code {}'.format(resp.status))
        return

    resp = json.loads(resp.read())

    error = resp.get('error')
    if error is not None:
        print('server responded with error message {}'.format(error))
        return

    result = resp.get('result')
    server_ephemeral_public_key = base64.b64decode(result.get('eph_pub_key'))

    #calculate the shared key
    r = libgksl.fhmqv_ecp_calc_shared((c_ubyte * FHMQV_ECP_SHARED_KEY_LENGTH).from_buffer(shared_key),
                                      client_static_private_key,
                                      client_static_public_key,
                                      (c_ubyte * FHMQV_ECP_PRIVATE_KEY_LENGTH).from_buffer(client_ephemeral_private_key),
                                      (c_ubyte * FHMQV_ECP_PUBLIC_KEY_LENGTH).from_buffer(client_ephemeral_public_key),
                                      server_static_public_key,
                                      server_ephemeral_public_key,
                                      1)

    session_id = base64.b64decode(result.get('session_id'))
    aes_key = hashlib.sha256(shared_key).digest()

    print("Session ID:")
    print(hexstr(session_id))
    print("Session key calculated:")
    print(hexstr(aes_key))


    #############################################################################################################
    ### APP-LEVEL AES-GCM ENCRYPTED COMMUNICATION

    iv = os.urandom(IV_SIZE)
    atag = bytearray(ATAG_SIZE)

    app_request = b'Hello Alice!'
    app_request_len = len(app_request)
    encrypted_request = bytearray(app_request_len)

    if libgksl.aes256_gcm_enc(aes_key,
                              iv, IV_SIZE,
                              (c_ubyte * ATAG_SIZE).from_buffer(atag), ATAG_SIZE,
                              session_id, SESSION_ID_SIZE,
                              app_request,
                              (c_ubyte * app_request_len).from_buffer(encrypted_request), app_request_len) != 0:
        print "aes-gcm encoding failed"
        return


    request = {'id': 0,
               'method': 'app_request',
               'params': {
                   'session_id': base64.b64encode(session_id),
                   'iv': base64.b64encode(iv),
                   'atag':  base64.b64encode(atag),
                   'request': base64.b64encode(encrypted_request)
               }}

    #The AES-GCM encrypted request is sent to the server
    conn.request('POST', '/', json.dumps(request))

    #TODO: Handle server reply


    #############################################################################################################
    ### CLOSE CONNECTION

    conn.close()

if __name__ == "__main__":
    main()